import os
import cv2
import numpy as np
import uuid
import subprocess
import time
import traceback
import sys
from PyQt6.QtCore import QObject, pyqtSignal, QThread

class ProcessWorker(QThread):
    """图像处理工作线程"""
    progress_updated = pyqtSignal(int, str)  # 进度信号：百分比，当前状态描述
    process_finished = pyqtSignal(object)    # 处理完成信号：处理结果
    process_error = pyqtSignal(str)          # 处理错误信号：错误信息
    
    def __init__(self, processor, image_path, params):
        super().__init__()
        self.processor = processor
        self.image_path = image_path
        self.params = params
    
    def run(self):
        try:
            # 通知开始处理
            self.progress_updated.emit(0, f"开始处理图像: {os.path.basename(self.image_path)}")
            
            # 读取图像 - 使用numpy处理中文路径问题
            try:
                # 使用二进制模式打开文件
                with open(self.image_path, 'rb') as f:
                    img_array = np.frombuffer(f.read(), np.uint8)
                    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
                if img is None:
                    raise ValueError(f"无法读取图像: {self.image_path}")
            except Exception as e:
                raise ValueError(f"无法读取图像: {self.image_path}. 错误: {str(e)}")
            
            # 通知进度
            self.progress_updated.emit(10, "准备处理参数")
            
            # 获取参数
            model_type = self.params.get('model_type', 'anime_style_art_rgb')
            noise_level = self.params.get('noise_level', 1)
            scale_factor = self.params.get('scale', 2)
            tta_mode = self.params.get('tta_mode', False)
            processor = self.params.get('processor', 'gpu')
            
            # 输出处理信息
            self.progress_updated.emit(20, f"使用{model_type}模型处理中...")
            
            # 获取Waifu2x-Caffe模型路径
            model_path = os.path.join(
                os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                'models',
                'waifu2x-caffe'
            )
            
            # 检查模型路径是否存在
            if not os.path.exists(model_path):
                raise ValueError("Waifu2x-Caffe模型路径不存在")
            
            # 生成临时文件名
            timestamp = int(time.time())
            random_id = uuid.uuid4().hex[:8]
            input_file = os.path.join(os.environ.get('TEMP', '/tmp'), f"waifu2x_in_{timestamp}_{random_id}.png")
            output_file = os.path.join(os.environ.get('TEMP', '/tmp'), f"waifu2x_out_{timestamp}_{random_id}.png")
            
            # 保存输入图像 - 使用imencode解决中文路径问题
            _, img_encoded = cv2.imencode('.png', img)
            with open(input_file, 'wb') as f:
                f.write(img_encoded)
            
            self.progress_updated.emit(30, "保存临时文件完成")
            
            # 构造`waifu2x-caffe-console.exe`命令
            exe_path = os.path.join(model_path, "waifu2x-caffe-console.exe")
            
            if not os.path.exists(exe_path):
                raise ValueError("waifu2x-caffe-console.exe不存在")
            
            # 检查模型文件夹
            model_dir = os.path.join(model_path, "models")
            if not os.path.exists(os.path.join(model_dir, model_type)):
                model_type = "anime_style_art_rgb"  # 默认回退到动漫风格
                self.progress_updated.emit(40, f"找不到指定模型，使用默认模型: {model_type}")
            
            # 构造命令参数
            cmd = [
                exe_path,
                "-i", input_file,
                "-o", output_file,
                "-m", "noise_scale",  # 降噪并放大
                "-s", str(scale_factor),  # 放大倍数
                "-n", str(noise_level),  # 降噪等级
                "-p", processor  # 处理器选择 (cpu/gpu)
            ]
            
            # 如果启用TTA模式
            if tta_mode:
                cmd.append("-x")
            
            # 执行命令 - 使用二进制模式避免编码问题
            self.progress_updated.emit(50, "正在处理图像...")
            # 使用二进制模式，不进行文本解码
            result = subprocess.run(cmd, check=True, capture_output=True)
            self.progress_updated.emit(80, "图像处理完成")
            
            # 检查输出文件是否存在
            if not os.path.exists(output_file):
                raise ValueError("处理后的图像文件不存在")
            
            # 读取处理后的图像 - 使用二进制读取解决中文路径问题
            with open(output_file, 'rb') as f:
                img_array = np.frombuffer(f.read(), np.uint8)
                result_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
                
            if result_img is None:
                raise ValueError("无法读取处理后的图像")
            
            # 清理临时文件
            try:
                os.remove(input_file)
                os.remove(output_file)
            except Exception as e:
                print(f"清理临时文件失败: {str(e)}")
            
            # 创建结果数据
            result_data = {
                'image': result_img,  # 处理后的图像
                'original_path': self.image_path,  # 原始图像路径
                'scale_factor': scale_factor,  # 放大倍数
                'model_type': model_type,  # 使用的模型类型
                'noise_level': noise_level  # 降噪等级
            }
            
            # 通知处理完成
            self.progress_updated.emit(100, "处理完成")
            self.process_finished.emit(result_data)
        
        except Exception as e:
            error_msg = f"处理图像失败: {str(e)}"
            print(error_msg)
            traceback.print_exc()
            self.process_error.emit(error_msg)

class BatchProcessWorker(QThread):
    """批量图像处理工作线程"""
    progress_updated = pyqtSignal(int, str)  # 进度信号：总体百分比，当前状态描述
    file_progress_updated = pyqtSignal(str, int)  # 单个文件进度信号：文件名，百分比
    process_finished = pyqtSignal(list)  # 处理完成信号：处理结果列表
    process_error = pyqtSignal(str)      # 处理错误信号：错误信息
    
    def __init__(self, processor, image_paths, output_dir, params):
        super().__init__()
        self.processor = processor
        self.image_paths = image_paths
        self.output_dir = output_dir
        self.params = params
    
    def run(self):
        try:
            results = []
            total_files = len(self.image_paths)
            
            # 获取参数
            model_type = self.params.get('model_type', 'anime_style_art_rgb')
            noise_level = self.params.get('noise_level', 1)
            scale_factor = self.params.get('scale', 2)
            tta_mode = self.params.get('tta_mode', False)
            processor = self.params.get('processor', 'gpu')
            
            for i, image_path in enumerate(self.image_paths):
                current_file = os.path.basename(image_path)
                self.progress_updated.emit(
                    int((i / total_files) * 100),
                    f"正在处理 [{i+1}/{total_files}]: {current_file}"
                )
                
                try:
                    # 处理单个图像 - 使用二进制读取解决中文路径问题
                    try:
                        with open(image_path, 'rb') as f:
                            img_array = np.frombuffer(f.read(), np.uint8)
                            img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
                        if img is None:
                            print(f"无法读取图像: {image_path}")
                            continue
                    except Exception as e:
                        print(f"无法读取图像: {image_path}. 错误: {str(e)}")
                        continue
                    
                    # 更新文件进度
                    self.file_progress_updated.emit(current_file, 20)
                    
                    # 获取Waifu2x-Caffe模型路径
                    model_path = os.path.join(
                        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                        'models',
                        'waifu2x-caffe'
                    )
                    
                    # 生成临时文件名
                    timestamp = int(time.time())
                    random_id = uuid.uuid4().hex[:8]
                    input_file = os.path.join(os.environ.get('TEMP', '/tmp'), f"waifu2x_in_{timestamp}_{random_id}.png")
                    output_file = os.path.join(os.environ.get('TEMP', '/tmp'), f"waifu2x_out_{timestamp}_{random_id}.png")
                    
                    # 保存输入图像 - 使用imencode解决中文路径问题
                    _, img_encoded = cv2.imencode('.png', img)
                    with open(input_file, 'wb') as f:
                        f.write(img_encoded)
                        
                    self.file_progress_updated.emit(current_file, 30)
                    
                    # 构造命令参数
                    exe_path = os.path.join(model_path, "waifu2x-caffe-console.exe")
                    cmd = [
                        exe_path,
                        "-i", input_file,
                        "-o", output_file,
                        "-m", "noise_scale",
                        "-s", str(scale_factor),
                        "-n", str(noise_level),
                        "-p", processor
                    ]
                    
                    # 如果启用TTA模式
                    if tta_mode:
                        cmd.append("-x")
                    
                    # 执行命令 - 使用二进制模式避免编码问题
                    self.file_progress_updated.emit(current_file, 50)
                    # 使用二进制模式，不进行文本解码
                    result = subprocess.run(cmd, check=True, capture_output=True)
                    self.file_progress_updated.emit(current_file, 80)
                    
                    # 检查输出文件是否存在
                    if not os.path.exists(output_file):
                        print(f"处理后的图像文件不存在: {output_file}")
                        continue
                    
                    # 读取处理后的图像 - 使用二进制读取解决中文路径问题
                    with open(output_file, 'rb') as f:
                        img_array = np.frombuffer(f.read(), np.uint8)
                        result_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
                    
                    if result_img is None:
                        print(f"无法读取处理后的图像: {output_file}")
                        continue
                    
                    # 保存到输出目录或原目录
                    save_dir = self.output_dir if self.output_dir else os.path.dirname(image_path)
                    base_name, ext = os.path.splitext(os.path.basename(image_path))
                    output_path = os.path.join(save_dir, f"{base_name}_{scale_factor}x{ext}")
                    
                    # 保存结果图像 - 使用imencode解决中文路径问题
                    _, img_encoded = cv2.imencode(ext, result_img)
                    with open(output_path, 'wb') as f:
                        f.write(img_encoded)
                    
                    # 添加到结果列表
                    results.append(output_path)
                    
                    # 清理临时文件
                    try:
                        os.remove(input_file)
                        os.remove(output_file)
                    except Exception as e:
                        print(f"清理临时文件失败: {str(e)}")
                    
                    self.file_progress_updated.emit(current_file, 100)
                    
                except Exception as e:
                    print(f"处理文件 {image_path} 失败: {str(e)}")
                    traceback.print_exc()
            
            # 处理完成
            self.progress_updated.emit(100, f"批量处理完成，成功处理 {len(results)} / {total_files} 个文件")
            self.process_finished.emit(results)
        
        except Exception as e:
            error_msg = f"批量处理失败: {str(e)}"
            print(error_msg)
            traceback.print_exc()
            self.process_error.emit(error_msg)

class ImageProcessor(QObject):
    """图像处理器类，负责图像的超分辨率处理"""
    
    progress_updated = pyqtSignal(int, str)  # 进度信号：百分比，当前状态描述
    file_progress_updated = pyqtSignal(str, int)  # 单个文件进度信号：文件名，百分比
    process_finished = pyqtSignal(object)  # 处理完成信号：处理结果
    process_error = pyqtSignal(str)  # 处理错误信号：错误信息
    batch_finished = pyqtSignal(list)  # 批量处理完成信号：结果文件列表
    
    def __init__(self):
        super().__init__()
        self.current_worker = None
        self.batch_worker = None
        
        # 临时目录
        self.temp_dir = os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            'temp'
        )
        if not os.path.exists(self.temp_dir):
            os.makedirs(self.temp_dir)
        
        # 检查Waifu2x-Caffe模型
        self.model_path = os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            'models',
            'waifu2x-caffe'
        )
        
        # 确保模型目录存在
        if not os.path.exists(self.model_path):
            os.makedirs(self.model_path, exist_ok=True)
    
    def is_model_available(self):
        """检查Waifu2x-Caffe模型是否可用"""
        # 检查可执行文件
        exe_path = os.path.join(self.model_path, "waifu2x-caffe-console.exe")
        return os.path.exists(exe_path)
    
    def process_image(self, image_path, params):
        """处理单个图像
        
        Args:
            image_path: 输入图像路径
            params: 处理参数字典
        """
        # 创建并启动工作线程
        self.current_worker = ProcessWorker(self, image_path, params)
        self.current_worker.progress_updated.connect(self.progress_updated)
        self.current_worker.process_finished.connect(self.process_finished)
        self.current_worker.process_error.connect(self.process_error)
        self.current_worker.start()
    
    def batch_process(self, image_paths, output_dir, params):
        """批量处理图像
        
        Args:
            image_paths: 输入图像路径列表
            output_dir: 输出目录（如果为None，则使用原图像目录）
            params: 处理参数字典
        """
        # 创建并启动批量处理线程
        self.batch_worker = BatchProcessWorker(self, image_paths, output_dir, params)
        self.batch_worker.progress_updated.connect(self.progress_updated)
        self.batch_worker.file_progress_updated.connect(self.file_progress_updated)
        self.batch_worker.process_finished.connect(self.batch_finished)
        self.batch_worker.process_error.connect(self.process_error)
        self.batch_worker.start()
